import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.classic_strategy.turtle_strategy import TurtleTrade


register_matplotlib_converters()


def main(codes, start):
    # print(codes.head())
    data_today = ts.get_today_all()
    data_today.to_csv('today.csv')
    # data_today = pd.read_csv('today.csv')
    data_today = data_today[['code', 'open', 'high', 'low', 'trade', 'volume', 'amount']]
    data_today = data_today.rename(columns={"trade": "close"})
    data_today['date'] = datetime.date.today()
    data_today.set_index(['date', 'code'], inplace=True)
    data = QA.QA_fetch_stock_day_adv(codes['code'].to_list(), start, end=str(datetime.date.today() + datetime.timedelta(days=-1)))
    data.data = data.data.append(data_today)
    # data.data.to_csv('all.csv')
    # print('saved')
    # data = data.to_qfq()
    ind = data.add_func(QA.QA_indicator_UP_SERIES, N=30)
    for _, code in codes.iterrows():
        df = ind.xs(code['code'], level='code', axis=0, drop_level=True)
        # df.to_csv('601933_zlxc.csv')
        ind_last_last_bar = df.iloc[-3]
        ind_last_bar = df.iloc[-2]
        ind_current_bar = df.iloc[-1]
        if ind_current_bar.B > 0:
            print('code：{}---{}---时间{}---CLOSE{}'.format(
                code['code'], code['name'], ind_current_bar.name, ind_current_bar.CLOSE))

            figure = plt.figure()
            plt.title('{}---'.format(code['code']))
            axes1 = figure.add_subplot(411)
            axes2 = figure.add_subplot(412)
            axes3 = figure.add_subplot(413)
            axes4 = figure.add_subplot(414)
            axes1.plot(df.index, df['CLOSE'])
            axes2.plot(df.index, df['VOLUME'])
            axes3.plot(df.index, df[['MA5', 'MA20', 'MA30']])
            axes4.plot(df.index, df[['B']])
            axes4.plot(df.index, df[['S']])

            left, right = axes2.get_xlim()
            axes2.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
            multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)

            plt.show()

    # data.to_csv('002581.csv')
    # print(data.index)


if __name__ == '__main__':
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    data = data[data.symbol.str.startswith('00') | data.symbol.str.startswith('60') | data.symbol.str.startswith('30')]
    # data = data[data.symbol.str.startswith('60088')]
    data = data[data.list_date < '20190101']
    all = data.rename(columns={'symbol': 'code'})
    codes = all[~all.name.str.contains('ST')]
    # codes = codes.iloc[0: 1]
    # hs300 = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\AI\\hs300.csv')
    # hs300 = ts.get_hs300s()
    # codes = codes['code'].to_list()
    # codes = ['002027']
    # print(codes)
    start = '2020-08-01'
    main(codes, start)
